package cn.lingyangwl.agile.gateway.filter;

import cn.lingyangwl.agile.gateway.service.FileService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;


/**
 * 对请求体进行修改‘
 * @author shenguangyang
 */
@Component
@Slf4j
public class RequestBodyUpdateGlobalFilter implements GlobalFilter, Ordered {
    @Resource
    private FileService fileService;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        AtomicReference<Mono<Void>> mono = new AtomicReference<>(chain.filter(exchange));
        MediaType contentType = exchange.getRequest().getHeaders().getContentType();
        if (Objects.nonNull(contentType) && Objects.nonNull(exchange.getRequest().getMethod())
                && exchange.getRequest().getMethod().equals(HttpMethod.POST)) {
            if (MediaType.APPLICATION_JSON.isCompatibleWith(contentType)) {
                // json 请求体处理
                mono.set(this.transferBody(exchange, chain));
            }
        }
        return mono.get();
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 100;
    }

    /**
     * 修改原请求体内容
     */
    @SuppressWarnings("all")
    private Mono<Void> transferBody(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        ServerRequest serverRequest = ServerRequest.create(exchange, HandlerStrategies.withDefaults().messageReaders());
        Mono modifiedBody = serverRequest.bodyToMono(String.class).flatMap(oldBody -> {
            // 这里也可以将修改后的请求体覆盖到GatewayContext缓存中，这里没有覆盖是因为想要保留最原始的请求体内容
            String newBody = fileService.realToVirtualBaseUrl(request, oldBody);
            return Mono.just(newBody);
        });
        BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());
        headers.remove(HttpHeaders.CONTENT_LENGTH);
        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
        Mono<Void> mono = bodyInserter.insert(outputMessage, new BodyInserterContext())
                .then(Mono.defer(() -> {
                    ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(
                            exchange.getRequest()) {
                        @Override
                        public HttpHeaders getHeaders() {
                            long contentLength = headers.getContentLength();
                            HttpHeaders httpHeaders = new HttpHeaders();
                            httpHeaders.putAll(super.getHeaders());
                            if (contentLength > 0) {
                                httpHeaders.setContentLength(contentLength);
                            } else {
                                httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                            }
                            return httpHeaders;
                        }

                        @Override
                        public Flux<DataBuffer> getBody() {
                            return outputMessage.getBody();
                        }
                    };
                    return chain.filter(exchange.mutate().request(decorator).build());
                }));
        return mono;
    }
}
